from matplotlib import pyplot as plt
import numpy as np
import os
from pyspark import SparkContext
from pyspark.mllib.classification import SVMWithSGD
from pyspark.mllib.linalg import DenseVector
from pyspark.mllib.regression import LabeledPoint

os.environ['PYSPARK_PYTHON'] = "python3"

red_points = [(1.3, 2.53), (1.2, 2.03), (1.73, 2.0), (2.33, 2.23), (2.37, 3.07), (1.73, 3.47), (1.2, 3.47), (0.83, 2.8), (1.27, 3.07), (2.0, 2.9), (2.0, 2.3), (1.6, 2.43), (1.6, 2.83), (1.9, 3.03), (2.1, 3.63), (2.17, 3.23), (2.63, 2.33), (2.57, 2.47), (2.3, 2.67), (0.97, 2.5), (0.93, 2.1), (1.93, 2.43), (1.97, 2.7), (2.57, 2.8), (1.83, 3.27), (1.5, 2.87), (1.07, 2.63), (1.37, 2.17), (2.2, 2.53)]
blue_points = [(3.53, 0.5), (4.03, 0.47), (4.53, 0.43), (4.57, 1.17), (3.93, 1.27), (3.23, 1.1), (3.3, 0.93), (3.67, 1.67), (4.17, 1.97), (4.37, 2.43), (4.13, 1.63), (4.6, 1.87), (4.5, 1.23), (4.0, 1.03), (3.67, 1.53), (3.43, 1.5), (3.5, 1.1), (3.5, 1.43), (3.6, 1.4), (3.73, 1.17), (3.73, 0.8), (3.47, 0.73), (3.97, 0.93), (3.83, 1.0), (4.0, 0.7), (4.53, 0.67), (4.37, 1.33), (4.03, 1.2), (4.17, 1.33), (4.43, 1.7), (4.4, 2.2), (4.0, 2.13), (3.97, 1.73), (3.9, 1.9), (3.97, 1.67), (4.17, 1.7), (4.3, 1.9), (4.3, 0.93), (4.27, 1.03), (4.33, 1.0), (4.33, 0.5), (4.17, 0.8), (4.53, 0.77), (4.83, 1.07), (4.87, 1.7), (4.9, 2.17), (4.5, 2.23), (4.7, 2.3), (4.7, 2.13), (4.73, 1.57), (4.7, 1.23)]


sc = SparkContext.getOrCreate()

labeled_points = [LabeledPoint(0, DenseVector((x, y))) for x, y in red_points] + [LabeledPoint(1, DenseVector((x, y))) for x, y in blue_points]

model = SVMWithSGD.train(sc.parallelize(labeled_points),iterations=200)

plt.scatter([x for x, _ in red_points], [y for _, y in red_points], color="r")
plt.scatter([x for x, _ in blue_points], [y for _, y in blue_points], color="b")

x = np.arange(0, 5, 0.1)

# print(model.weights, model.intercept)
plt.ylim(0, 5)
plt.plot(x, (model.weights[0])* x + model.weights[1] , color="g")
# plt.plot(x, model.weights[1] * x, color="g")

plt.show()

